import streamlit as st
from spider import init_driver, wait_for_login, get_jobs
from config import zone_dicts, exp_list, scale_list
import pandas as pd
import os
from selenium.common.exceptions import WebDriverException
import traceback

st.set_page_config(
    page_title="BOSS直聘爬虫",
    page_icon="💼",
    layout="wide",
    initial_sidebar_state="expanded"
)

st.title("💼 BS直聘职位爬虫")
st.markdown("---")

# 初始化 session state
if "driver" not in st.session_state:
    st.session_state.driver = None
if "logged_in" not in st.session_state:
    st.session_state.logged_in = False

# 侧边栏配置
with st.sidebar:
    st.header("🔍 搜索配置")
    
    # 搜索条件
    city = st.selectbox("选择城市", list(zone_dicts.keys()))
    keyword = st.text_input("输入职位关键词", value="算法工程师")
    exp = st.selectbox("工作经验", exp_list)
    scale = st.selectbox("公司规模", scale_list)
    limit = st.slider("爬取数量", min_value=10, max_value=100, value=20, step=10)
    
    # 登录按钮
    if not st.session_state.logged_in and st.button("🔑 开始扫码登录", use_container_width=True):
        st.session_state["start"] = True
        
    # 显示配置信息
    st.markdown("---")
    st.markdown("### ⚙️ 当前配置")
    st.write(f"- 城市: {city}")
    st.write(f"- 关键词: {keyword}")
    st.write(f"- 工作经验: {exp}")
    st.write(f"- 公司规模: {scale}")
    st.write(f"- 爬取数量: {limit}")

# 主页面
if st.session_state.get("start"):
    try:
        if st.session_state.driver is None:
            with st.spinner("⏳ 正在启动浏览器..."):
                try:
                    st.session_state.driver = init_driver()
                except Exception as e:
                    st.error(f"❌ 浏览器启动失败: {str(e)}")
                    st.error("请确保您的系统已安装 Chrome 浏览器")
                    st.stop()
        
        if not st.session_state.logged_in:
            with st.spinner("⏳ 等待扫码登录...（请使用 BOSS直聘APP 扫描二维码登录）"):
                try:
                    wait_for_login(st.session_state.driver)
                    st.session_state.logged_in = True
                    st.success("✅ 登录成功！")
                    st.rerun()
                except Exception as e:
                    st.error(f"❌ 登录失败: {str(e)}")
                    if st.session_state.driver:
                        st.session_state.driver.quit()
                        st.session_state.driver = None
                    st.session_state.logged_in = False
                    st.stop()
        
        if st.session_state.logged_in:
            if st.button("🚀 开始爬取数据", use_container_width=True):
                try:
                    with st.spinner(f"🔄 正在爬取数据...（预计需要 {limit*3} 秒）"):
                        filename = get_jobs(
                            driver=st.session_state.driver,
                            keyword=keyword,
                            city_code=zone_dicts[city]['code'],
                            exp=exp,
                            scale=scale,
                            limit=limit
                        )
                    
                    # 读取并显示数据
                    df = pd.read_csv(filename)
                    st.success(f"✅ 成功爬取 {len(df)} 条数据！")
                    
                    # 数据预览
                    st.markdown("### 📊 数据预览")
                    st.dataframe(df, use_container_width=True)
                    
                    # 下载按钮
                    with open(filename, "rb") as f:
                        st.download_button(
                            "📥 下载数据(CSV)",
                            f,
                            file_name=filename.split("/")[-1],
                            mime="text/csv",
                            use_container_width=True
                        )
                        
                except WebDriverException as e:
                    st.error(f"❌ 浏览器操作错误: {str(e)}")
                    st.error("请尝试重新启动应用")
                    st.session_state.logged_in = False
                except Exception as e:
                    st.error(f"❌ 发生错误: {str(e)}")
                    st.error(f"错误详情:\n```\n{traceback.format_exc()}\n```")
                finally:
                    if st.session_state.driver:
                        st.session_state.driver.quit()
                        st.session_state.driver = None
                    st.session_state.logged_in = False
                    
    except Exception as e:
        st.error(f"❌ 发生未知错误: {str(e)}")
        if st.session_state.driver:
            st.session_state.driver.quit()
            st.session_state.driver = None
        st.session_state.logged_in = False
else:
    st.info("👈 请在左侧设置搜索条件并点击登录按钮开始")
    
# 页脚
st.markdown("---")
st.markdown("### 📝 使用说明")
st.markdown("""
1. 在左侧设置搜索条件
2. 点击"开始扫码登录"按钮
3. 使用BS直聘APP扫描二维码登录（无需手动确认）
4. 登录成功后，点击"开始爬取数据"按钮
5. 等待数据爬取完成，可以预览和下载数据

**注意事项：**
- 确保系统已安装最新版本的 Chrome 浏览器
- 如果遇到错误，请尝试刷新页面重新开始
- 建议每次爬取数量不要太大，以避免被反爬
- 爬取过程中请勿关闭浏览器窗口
- 每个职位详情页面会有 2 秒延迟，请耐心等待
""")
